{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "import IPython\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import soundfile as sf\n",
    "from tqdm import tqdm\n",
    "\n",
    "from nara_wpe.wpe import wpe\n",
    "from nara_wpe.wpe import get_power\n",
    "from nara_wpe.utils import stft, istft, get_stft_center_frequencies\n",
    "from nara_wpe import project_root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stft_options = dict(size=512, shift=128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Minimal example with random data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def aquire_audio_data():\n",
    "    D, T = 4, 10000\n",
    "    y = np.random.normal(size=(D, T))\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = aquire_audio_data()\n",
    "Y = stft(y, **stft_options)\n",
    "Y = Y.transpose(2, 0, 1)\n",
    "\n",
    "Z = wpe(Y)\n",
    "z_np = istft(Z.transpose(1, 2, 0), size=stft_options['size'], shift=stft_options['shift'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example with real audio recordings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "WPE estimates a filter to predict the current reverberation tail frame from K time frames which lie 3 (delay) time frames in the past. This frame (reverberation tail) is then subtracted from the observed signal.\n",
    "\n",
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "channels = 8\n",
    "sampling_rate = 16000\n",
    "delay = 3\n",
    "iterations = 5\n",
    "taps = 10\n",
    "alpha=0.9999"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Audio data\n",
    "Shape: (channels, frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_template = 'AMI_WSJ20-Array1-{}_T10c0201.wav'\n",
    "signal_list = [\n",
    "    sf.read(str(project_root / 'data' / file_template.format(d + 1)))[0]\n",
    "    for d in range(channels)\n",
    "]\n",
    "y = np.stack(signal_list, axis=0)\n",
    "IPython.display.Audio(y[0], rate=sampling_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### STFT\n",
    "A STFT is performed to obtain a Numpy array with shape (frequency bins, channels, frames)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = stft(y, **stft_options).transpose(2, 0, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iterative WPE\n",
    "The wpe function is fed with Y. Finally, an inverse STFT is performed to obtain a dereverberated result in time domain. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z = wpe(\n",
    "    Y,\n",
    "    taps=taps,\n",
    "    delay=delay,\n",
    "    iterations=iterations,\n",
    "    statistics_mode='full'\n",
    ").transpose(1, 2, 0)\n",
    "z = istft(Z, size=stft_options['size'], shift=stft_options['shift'])\n",
    "IPython.display.Audio(z[0], rate=sampling_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Power spectrum \n",
    "Before and after applying WPE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, [ax1, ax2] = plt.subplots(1, 2, figsize=(20, 10))\n",
    "im1 = ax1.imshow(20 * np.log10(np.abs(Y[ :, 0, 200:400])), origin='lower')\n",
    "ax1.set_xlabel('frames')\n",
    "_ = ax1.set_title('reverberated')\n",
    "im2 = ax2.imshow(20 * np.log10(np.abs(Z[0, 200:400, :])).T, origin='lower', vmin=-120, vmax=0)\n",
    "ax2.set_xlabel('frames')\n",
    "_ = ax2.set_title('dereverberated')\n",
    "cb = fig.colorbar(im2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
